import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QLabel
from PyQt5.QtGui import QPainter, QPen, QPolygon, QColor, QFont
from PyQt5.QtCore import Qt, QPoint

# 状态机类
class StateMachine:
    def __init__(self):
        self.state = 0
        self.state_names = {
            0: "Start (0)",
            1: "Not ready \n to switch on (1)",
            2: "Switch on \n disabled (2)",
            3: "Ready to \n switch on (3)",
            4: "Switched on (4)",
            5: "Operation \n enabled (5)",
            8: "Fault reaction \n active (8)",
            9: "Fault (9,14)",
            11: "Quick stop (11)",
            13: "Fault reaction \n active (13)",
        }

    def transition(self, condition):
        if self.state == 0:
            self.state = 1
        elif self.state == 1:
            if condition == "switch_on":
                self.state = 2
            else:
                self.state = 9
        elif self.state == 2:
            if condition == "switch_on":
                self.state = 3
        elif self.state == 3:
            if condition == "switch_on":
                self.state = 4
            elif condition == "disable":
                self.state = 2
            else:
                self.state = 2  # 原 3 -> 6 改为 3 -> 2
        elif self.state == 4:
            if condition == "enable_operation":
                self.state = 5
            elif condition == "power_off_or_reset":
                self.state = 1
            elif condition == "disable":
                self.state = 3
            else:
                self.state = 8
        elif self.state == 5:
            if condition == "disable":
                self.state = 4
            elif condition == "quick_stop":
                self.state = 11
            elif condition == "disable":
                self.state = 2
            else:
                self.state = 9
        elif self.state == 8:
            self.state = 9
        elif self.state == 9:
            if condition == "power_off_or_reset":
                self.state = 0
            elif condition == "switch_on":
                self.state = 2
        elif self.state == 11:
            self.state = 2
        elif self.state == 13:
            self.state = 9

# 状态机窗口类
class StateMachineWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("State Machine: Trigger States")
        self.setGeometry(100, 100, 1200, 800)

        # 状态机实例
        self.sm = StateMachine()

        # 状态控件位置和大小
        self.state_positions = {
            0: (550, 50, 200, 50),   # Start (0)
            1: (550, 150, 200, 50),  # Not ready to switch on (1)
            2: (550, 250, 200, 50),  # Switch on disabled (2)
            3: (550, 350, 200, 50),  # Ready to switch on (3)
            4: (550, 450, 200, 50),  # Switched on (4)
            5: (550, 550, 200, 50),  # Operation enabled (5)
            8: (800, 450, 200, 50),  # Fault reaction active (8)
            9: (950, 550, 200, 50),  # Fault (9,14)
            11: (200, 550, 200, 50), # Quick stop (11)
            13: (950, 250, 200, 50), # Fault reaction active (13)
        }

        # 按键提示标签
        instructions = [
            "Key Instructions:",
            "1: Switch on",
            "2: Enable operation",
            "3: Disable",
            "4: Quick stop",
            "5: Power off or reset"
        ]
        for i, instruction in enumerate(instructions):
            label = QLabel(instruction, self)
            label.setFont(QFont("Arial", 12))
            label.move(50, 50 + i * 30)

    def keyPressEvent(self, event):
        # 按键映射
        key_mappings = {
            Qt.Key_1: "switch_on",
            Qt.Key_2: "enable_operation",
            Qt.Key_3: "disable",
            Qt.Key_4: "quick_stop",
            Qt.Key_5: "power_off_or_reset"
        }
        condition = key_mappings.get(event.key())
        if condition:
            self.sm.transition(condition)
            self.update()  # 触发重绘

    def paintEvent(self, event):
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)

        # 绘制状态框
        for state, (x, y, w, h) in self.state_positions.items():
            # 根据状态设置颜色
            if state == self.sm.state:
                color = QColor(0, 255, 0)  # 绿色
            elif state in [9]:
                color = QColor(255, 0, 0)  # 红色
            else:
                color = QColor(200, 200, 200)  # 灰色

            # 绘制矩形
            painter.setPen(QPen(Qt.black, 2))
            painter.setBrush(color)
            painter.drawRect(x, y, w, h)

            # 绘制状态名称
            painter.setPen(Qt.black)
            painter.setFont(QFont("Arial", 10))
            painter.drawText(x, y, w, h, Qt.AlignCenter, self.sm.state_names[state])

        # 绘制箭头和标签
        arrows = [
            # 0 -> 1 (直接向下)
            ([((650, 100), (650, 150))], "", (0, 0), "down"),
            # 1 -> 2 (直接向下)
            ([((650, 200), (650, 250))], "SWITCH ON", (700, 225), "down"),
            # 2 -> 3 (直接向下)
            ([((650, 300), (650, 350))], "SWITCH ON", (700, 325), "down"),
            # 3 -> 4 (直接向下)
            ([((650, 400), (650, 450))], "SWITCH ON", (700, 425), "down"),
            # 4 -> 5 (直接向下)
            ([((650, 500), (650, 550))], "ENABLE OPERATION", (700, 525), "down"),
            # 5 -> 4 (向上)
            ([((650, 550), (650, 500))], "DISABLE", (700, 525), "up"),
            # 4 -> 3 (向上)
            ([((650, 450), (650, 400))], "DISABLE", (700, 425), "up"),
            # 3 -> 2 (向上)
            ([((650, 350), (650, 300))], "DISABLE", (700, 325), "up"),
            # 4 -> 1 (向上绕过 3, 2)
            ([((650, 450), (650, 400)), ((650, 400), (700, 400)), ((700, 400), (700, 100)), ((700, 100), (650, 100))], "POWER OFF", (750, 250), "up"),
            # 5 -> 2 (向上绕过 4, 3)
            ([((650, 550), (650, 500)), ((650, 500), (600, 500)), ((600, 500), (600, 200)), ((600, 200), (650, 200))], "DISABLE", (500, 350), "up"),
            # 4 -> 8 (直接向右)
            ([((750, 475), (800, 475))], "", (0, 0), "right"),
            # 5 -> 9 (向右绕过 8)
            ([((750, 575), (875, 575)), ((875, 575), (875, 600)), ((875, 600), (950, 575))], "", (0, 0), "right"),
            # 9 -> 8 (向上)
            ([((1050, 550), (1050, 450))], "", (0, 0), "up"),
            # 9 -> 0 (向上绕过 8)
            ([((1050, 550), (1050, 450)), ((1050, 450), (1100, 450)), ((1100, 450), (1100, 50)), ((1100, 50), (650, 50))], "POWER OFF", (1150, 250), "up"),
            # 9 -> 2 (向上绕过 8)
            ([((1050, 550), (1050, 450)), ((1050, 450), (900, 450)), ((900, 450), (900, 200)), ((900, 200), (650, 200))], "SWITCH ON", (850, 350), "up"),
            # 13 -> 9 (向下)
            ([((1050, 250), (1050, 550))], "", (0, 0), "down"),
            # 3 -> 2 (向左)
            ([((550, 375), (350, 375)), ((350, 375), (350, 300)), ((350, 300), (550, 300))], "", (0, 0), "left"),
            # 5 -> 11 (直接向左)
            ([((550, 575), (200, 575))], "QUICK STOP", (300, 625), "left"),
            # 11 -> 2 (直接向上)
            ([((300, 550), (300, 400)), ((300, 400), (650, 400)), ((650, 400), (650, 300))], "", (0, 0), "up"),
        ]

        painter.setPen(QPen(Qt.black, 2))
        for path, label, (label_x, label_y), direction in arrows:
            # 绘制多段路径
            for (start_x, start_y), (end_x, end_y) in path:
                painter.drawLine(start_x, start_y, end_x, end_y)

            # 绘制箭头头
            last_segment = path[-1]
            end_x, end_y = last_segment[1]
            if direction == "down":
                arrow = QPolygon([
                    QPoint(end_x, end_y),
                    QPoint(end_x - 5, end_y - 5),
                    QPoint(end_x + 5, end_y - 5)
                ])
            elif direction == "up":
                arrow = QPolygon([
                    QPoint(end_x, end_y),
                    QPoint(end_x - 5, end_y + 5),
                    QPoint(end_x + 5, end_y + 5)
                ])
            elif direction == "right":
                arrow = QPolygon([
                    QPoint(end_x, end_y),
                    QPoint(end_x - 5, end_y - 5),
                    QPoint(end_x - 5, end_y + 5)
                ])
            elif direction == "left":
                arrow = QPolygon([
                    QPoint(end_x, end_y),
                    QPoint(end_x + 5, end_y - 5),
                    QPoint(end_x + 5, end_y + 5)
                ])
            painter.setBrush(Qt.black)
            painter.drawPolygon(arrow)

            # 绘制标签
            if label:
                painter.setFont(QFont("Arial", 8))
                painter.drawText(label_x - 50, label_y - 5, 100, 20, Qt.AlignCenter, label)

# 主函数
if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = StateMachineWindow()
    window.show()
    sys.exit(app.exec_())